#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# # author : cypro666
# # date   : 2015.08.01
# # wrapper of KNN method in sklearn
import sys, math, json
import numpy as np
from threading import Thread
from sklearn import neighbors

from magic3.utils import Timer
from magic3.filesystem import *
add_sys_path(grand_dir(__file__))

from skt.base import MethodBase
from skt.utils import _normalize, make_options


class KNN(MethodBase):
    def __init__(self, parameters={}):
        super().__init__(parameters)

    @property
    def name(self): return 'KNN'

    def output(self, cls):
        self.Log('output')
        super().output()

        extra = {
            "proba" : [list(a) for a in cls.predict_proba(self._testing)]
        }

        fn = self.json_name()
        json.dump(extra, open(fn, 'w'), indent=4)

    def execute(self):
        self.read_input()
        self._normalize(False, '1')

        if self._param['num_neighbors'] < 2:
            self._param['num_neighbors'] = 2

        if self._param['leaf_size'] < 4:
            self._param['leaf_size'] = 4

        self.save_parameters()

        cls = neighbors.KNeighborsClassifier(algorithm=self._param['algorithm'],
                                             n_neighbors=self._param['num_neighbors'],
                                             leaf_size=self._param['leaf_size'],
                                             weights=self._param['weights'])

        cls.fit(self._train, self._label)
        self._results = cls.predict(self._testing)
        self.output(cls)

    def run(self, timeout):
        t = Thread(target=self.execute)
        t.start()
        t.join(timeout)
        if t.is_alive():
            self.Log('timeout!')
        self.Log('exit')


if __name__ == '__main__':
    opts = [('train_file', 'str', []),
            ('label_file', 'str', []),
            ('testing_file', 'str', []),
            ('results_file', 'str', []),
            ('log_file', 'str', []),
            ('num_neighbors', 'int', []),
            ('leaf_size', 'int', []),
            ('algorithm', 'choice', ['auto', 'ball_tree', 'kd_tree', 'brute']),
            ('weights', 'choice', ['uniform', 'distance'])]

    parameters = make_options(opts)

    timer = Timer()
    KNN(parameters).run(30000)

    if __debug__:
        print('elapsed:', timer)


